#include "llvm/ADT/VectorExtras.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/InstIterator.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/IRBuilder.h"
#include "llvm/Instructions.h"
#include "llvm/Module.h"
#include "llvm/Pass.h"
#include "llvm/Function.h"
#include "llvm/PassSupport.h"
#include "llvm/ValueSymbolTable.h"

//#include <cassert>
//#include <iterator>
//#include <vector>

using namespace llvm;

namespace {
class MPIBlock2NonBlock : public FunctionPass {
  Function *isend_func;
  Function *irecv_func;
  Function *wait_func;
  bool NeedTransform;
  void insertIsendOrIrecv(Function &F, CallInst* callInst,
      std::vector<CallInst*> &instRemoveList, bool SorR);
  Function *insertIsendOrIrecvDeclare(Module &M, bool SorR);
  Function *insertWaitDeclare(Module &M);

public:
  static char ID; // Pass identification, replacement for typeid
  MPIBlock2NonBlock() : FunctionPass(ID) {NeedTransform = false;}
  virtual bool doInitialization(Module &M);
  virtual bool runOnFunction(Function &F);

};

char MPIBlock2NonBlock::ID = 0;
static RegisterPass<MPIBlock2NonBlock> X("mpi-b2nb",
    "MPI Blocking to Non-Blocking Pass", true, false);

}

//const PassInfo *const MPIBlock2NonBlockID = &X;
//Pass *createMPIBlock2NonBlockPass() { return new MPIBlock2NonBlock(); }

bool MPIBlock2NonBlock::doInitialization(Module &M){
  bool Changed = false;
  //


  Function *send_func = M.getFunction(StringRef("MPI_Send"));
  Function *recv_func = M.getFunction(StringRef("MPI_Recv"));
  // check if preliminary transformation is needed.
  // search for MPI_Send and MPI_Recv function declaration.
  if(!send_func && !recv_func) {
    errs().write_escaped(StringRef("nothing to transform")) << '\n';
    return Changed;
  }

  NeedTransform = true;


  // add MPI_Isend declaration
  isend_func = M.getFunction(StringRef("MPI_Isend"));
  if(!isend_func){
    isend_func = insertIsendOrIrecvDeclare(M, true);
    Changed = true;
  }

  // add MPI_Irecv declaration
  irecv_func = M.getFunction(StringRef("MPI_Irecv"));
  if(!irecv_func){
    irecv_func = insertIsendOrIrecvDeclare(M, false);
    Changed = true;
  }

  // add MPI_Wait declaration
  wait_func = M.getFunction(StringRef("MPI_Wait"));
  if(!wait_func){
    wait_func = insertWaitDeclare(M);
    Changed = true;
  }

  return Changed;
}

bool MPIBlock2NonBlock::runOnFunction(Function &F) {
  bool Changed = false;
  if(!NeedTransform) return Changed;

  std::vector<CallInst*> instRemoveList;

  for(inst_iterator i = inst_begin(F), e = inst_end(F); i != e; ++i) {
    if (CallInst* callInst = dyn_cast<CallInst>(&*i)) {
      Function* func = callInst->getCalledFunction();
      if ((func->getName()).startswith(StringRef("MPI_Send"))) {
        insertIsendOrIrecv(F, callInst, instRemoveList, true);
      }
      if ((func->getName()).startswith(StringRef("MPI_Recv"))) {
        insertIsendOrIrecv(F, callInst, instRemoveList, false);
      }
    }
  }

  //erase MPI_Send & MPI_Recv instruction
  for(std::vector<CallInst*>::iterator i = instRemoveList.begin(),
      e = instRemoveList.end(); i != e; i++){
    (*i)->eraseFromParent();
  }

  Changed = true;
  return Changed;
}

Function *MPIBlock2NonBlock::insertIsendOrIrecvDeclare(Module &M, bool SorR) {
  /*** declare function MPI_Isend or MPI_Irecv prototype ***/
  std::vector<Type*> isend_7args = make_vector<Type*>(
      PointerType::getUnqual(IntegerType::getInt8Ty(M.getContext())),
      IntegerType::getInt32Ty(M.getContext()),
      IntegerType::getInt32Ty(M.getContext()),
      IntegerType::getInt32Ty(M.getContext()),
      IntegerType::getInt32Ty(M.getContext()),
      IntegerType::getInt32Ty(M.getContext()),
      PointerType::getUnqual(IntegerType::getInt32Ty(M.getContext())));

  FunctionType *isend_funcTy = FunctionType::get(
      /*Result=*/IntegerType::getInt32Ty(M.getContext()),
      /*Params=*/makeArrayRef(isend_7args),
      /*isVarArg=*/false);

  //SorR == true, send; SorR == false, recv
  Function *isend_func = Function::Create(
      /*Type=*/isend_funcTy,
      /*Linkage=*/GlobalValue::ExternalLinkage,
      /*Name=*/SorR?"MPI_Isend":"MPI_Irecv", &M); // (external, no body)

  return isend_func;
}

Function *MPIBlock2NonBlock::insertWaitDeclare(Module &M) {
  //  TypeSymbolTable &tst = M.getTypeSymbolTable();


  /*** declare function MPI_Wait prototype ***/
  std::vector<Type*> wait_2args = make_vector<Type*>(
      PointerType::getUnqual(IntegerType::getInt32Ty(M.getContext())),
      PointerType::getUnqual(M.getTypeByName("struct.MPI_Status")),
      0);

  FunctionType *wait_funcTy = FunctionType::get(
      /*Result=*/IntegerType::getInt32Ty(M.getContext()),
      /*Params=*/makeArrayRef(wait_2args),
      /*isVarArg=*/false);

  Function *wait_func = Function::Create(
      /*Type=*/wait_funcTy,
      /*Linkage=*/GlobalValue::ExternalLinkage,
      /*Name=*/"MPI_Wait", &M); // (external, no body)

  return wait_func;
}

void MPIBlock2NonBlock::insertIsendOrIrecv(Function &F,
    CallInst* callInst, std::vector<CallInst*> &instRemoveList, bool SorR) {
  Module &M = *(F.getParent());

  ValueSymbolTable &vst = F.getValueSymbolTable();

  /*** add request and status to entry BB of this function ***/
  //construct send_request & send_status name
  static int varNum;
  int t = ++varNum;
  Twine x1 = Twine(SorR?"send_status":"recv_status")+Twine(t);
  Twine x2 = Twine(SorR?"send_request":"recv_request")+Twine(t);

  IRBuilder<> Builder(&(F.front()).front());
  Builder.CreateAlloca(M.getTypeByName("struct.MPI_Status"), 0, x1);
  //const Type *StatusTy = M.getTypeByName("struct.MPI_Status");
  // Instruction *sstatusInst = new AllocaInst(StatusTy, x1, entry_inst);
  Type *ReqTy = PointerType::getUnqual(IntegerType::getInt32Ty(M.getContext()));
  Builder.CreateAlloca(ReqTy, 0, x2);
  //Instruction *sreqInst = new AllocaInst(ReqTy, x2, entry_inst);

  //create argument to the MPI_Wait call
  std::vector<Value*> wait_func_2params = make_vector<Value*>(
      vst.lookup(StringRef(x2.str())),
      vst.lookup(StringRef(x1.str())),
      0);

  DEBUG(errs() << "create call instruction of MPI_Wait.\n");
  //create call instruction of MPI_Wait
  //  CallInst* wait = CallInst::Create(
  //      wait_func,
  //      wait_func_2params.begin(),
  //      wait_func_2params.end(),
  //      "", callInst);

  Builder.SetInsertPoint(callInst);
  CallInst* wait = Builder.CreateCall2(wait_func, vst.lookup(StringRef(x2.str())),
      vst.lookup(StringRef(x1.str())));
  wait->setDoesNotThrow();//add nounwind to this call instruction
  DEBUG(errs() << "AFTER create call instruction of MPI_Wait.\n");

  //create argument to the MPI_Isend or MPI_Irecv call
  std::vector<Value*> func_7params;
  unsigned NumOperands = SorR? callInst->getNumOperands()-1: callInst->getNumOperands()-2 ;
  for(unsigned i = 0; i < NumOperands; i++){
    DEBUG(errs() << i << " " << callInst->getOperand(i)->getNameStr() << "\n");
    func_7params.push_back(callInst->getOperand(i));
  }
  func_7params.push_back(vst.lookup(StringRef(x2.str())));
  DEBUG(errs() << "AFTER create argument to the MPI_Isend or MPI_Irecv call.\n");

  //create call instruction of MPI_Isend
  CallInst* isend_irecv = CallInst::Create(
      SorR?isend_func:irecv_func,
          makeArrayRef(func_7params),
          "", wait);
  isend_irecv->setDoesNotThrow();//add nounwind to this call instruction
  DEBUG(errs() << "AFTER create call instruction of MPI_Isend.\n");


  //insert MPI_Send & MPI_Recv instruction handle into a list
  //erase MPI_Send & MPI_Recv instruction later
  instRemoveList.push_back(callInst);
}
